#!/usr/bin/python3
# filepath: /home/ma-user/work/TraceFramework/TraceKernelLaunch_tiling/scripts/gen_data.py
import os
import numpy as np

def trace(x):
    # x shape: [batch, rows, cols] -> 计算每个矩阵的trace
    result = []
    batch_size, rows, cols = x.shape
    for i in range(batch_size):
        # 提取对角线元素并求和
        diag_sum = 0
        min_dim = min(rows, cols)
        for j in range(min_dim):
            diag_sum += x[i, j, j]
        result.append(diag_sum)
    return np.array(result)

def gen_golden_data_simple():
    # 生成合适的矩阵维度
    batch_size = 15  # 矩阵数量
    matrix_rows = 777  # 每个矩阵的行数
    matrix_cols = 777  # 每个矩阵的列数
    shape = [batch_size, matrix_rows, matrix_cols]
    
    # 生成随机输入数据
    rng = np.random.default_rng()  # 不传入种子，确保每次不同
    input_x = rng.random(size=tuple(shape)).astype(np.float16)
    
    # 计算golden数据（使用float32精度计算再转回float16）
    input_float32 = input_x.astype(np.float32)
    golden = trace(input_float32).astype(np.float16)

    # tiling数据：[矩阵数量, 行数, 列数]
    tiling = np.array([batch_size, matrix_rows, matrix_cols], dtype=np.uint32)
    
    # 创建目录并保存文件
    os.system("mkdir -p input")
    os.system("mkdir -p output")
    
    input_x.tofile("./input/input_x.bin")
    tiling.tofile("./input/input_tiling.bin")
    golden.tofile("./output/golden.bin")

    print(f"Input shape: {input_x.shape}")
    print(f"Tiling values: {tiling}")
    print(f"Golden shape: {golden.shape}")
    print(f"Golden values: {golden}")

if __name__ == "__main__":
    gen_golden_data_simple()